import random
import torch
from d2l import torch as d2l
import numpy as np
from torch.utils import data
from torch.utils.data import DataLoader  # 由于使用data.Dataloader报错，这里直接导入
from torch import nn

# 生成数据集，x为1000个样本，w包含两个特征，手动给出w和b
def synthetic_data(w, b, num_examples):
    X = torch.normal(0, 1, (num_examples, len(w))) # 每个样本包含从正态分布中取样的两个特征
    y = torch.matmul(X, w)+b #线性函数y定义
    y += torch.normal(0, 0.01, y.shape) # 引入一个正态分布的噪声项生成标签
    return X, y.reshape(-1, 1) # 返回生成的X与标签y
true_w = torch.tensor([2, -3.4])
true_b = 4.2
features, labels = synthetic_data(true_w, true_b, 1000) # 最后调用函数生成需要的数据集
print('features:', features[0],'\nlabel:', labels[0])
d2l.set_figsize() # 画图
d2l.plt.scatter(features[:, (1)].detach().numpy(), labels.detach().numpy(), 1);

# 读取数据集
# 训练模型时要对数据集进行遍历，每次抽取一小批量样本，并使用它们来更新我们的模型
# ，我们定义一个data_iter函数，该函数接收批量大小、特征矩阵和标签向量作为输入，生成 大小为batch_size的小批量。每个小批量包含一组特征和标签
def data_iter(batch_size, features, labels):
    indices = list(range(0, len(features))) # 生成数组下标索引0~999，便于取随机batch
    # 这些样本是随机读取的，没有特定的顺序
    random.shuffle(indices)
    for i in range(0, len(features), batch_size): # 间隔10进行遍历
        batch_indices = torch.tensor(indices[i:min(i+batch_size, len(features))])
        # out：[990, 991, 992, 993, 994, 995, 996, 997, 998, 999]，只取最后一轮的i即990进行输出
    yield features[batch_indices], labels[batch_indices] # yield可以理解为特殊的return，处理大批量数据时常用
batch_size = 10
for X, y in data_iter(batch_size, features, labels): # 最后进行调用函数读取数据集
    print(X, '\n', y)
    break

# 初始化模型参数
# 过从均值为0、标准差为0.01的正态分布中采样随机数来初始化权重，并将偏置初始化为0，并打开自动求导
w = torch.normal(0, 0.01, size=(2, 1), requires_grad=True) # w为特征列向量,所以size（2，1）
b = torch.zeros(1, requires_grad=True)

# 定义模型
def linreg(X, w, b): # 线性回归模型
    return torch.matmul(X, w) + b # 利用了广播机制，因为xw是向量，b是标量

# 定义损失函数
def square_loss(y_hat, y):
    # 均方损失
    return (y_hat - y.reshape(y_hat.shape))**2 / 2 # 要将真实值y的形状转换为和预测值y_hat的形状相同

# 定义优化算法
def sgd(params, lr, batch_size):
    # 小批量随机梯度下降
    with torch.no_grad(): # no_gard()下包含的所用操作都不能进行求导，节省内存，防止更新参数时出错
        for params in params: # 注意所有参数都得更新，所以这里需要使用for循环
            params -= lr * params.grad / batch_size
            params.grad.zero_() # 每次都记得清零

# 训练模型,1、初始化参数 2、不断计算梯度并更新参数
lr = 0.03
num_epochs = 3
net = linreg
loss = square_loss # 初始化参数方便显示
# 关键的for循环
for epoch in range(num_epochs):
    for X, y in data_iter(batch_size, features, labels): # 读取小批量数据进行计算
        lossmin = loss(net(X, w, b), y) # 小批量数据的损失计算
        lossmin.sum().backward() # 反向传播计算关于[w,b]的梯度
        sgd([w, b], lr, batch_size) # 使用优化算法更新参数
    with torch.no_grad():
        train1 = loss(net(features, w, b), labels) # 最后在利用最新的参数进行计算训练损失，X，y可以用fea和label表示了
        print(f"epoch:{epoch+1},loss:{float(train1.mean())}")

# 比较训练得到的参数误差
print(f'w的估计误差: {true_w - w.reshape(true_w.shape)}')
print(f'b的估计误差: {true_b - b}')

# 调库实现线性回归，会简单一点
# 生成数据集，以上类似

# 读取数据集，使用api调库
def load_array(data_array, batch_size, is_train=True):
    # 构造一个pytorch数据迭代器
    dataset = data.TensorDataset(*data_array)
    return DataLoader(dataset, batch_size, shuffle=is_train) # 导入包错误，更换导入方式
data_iter = load_array((features, labels), batch_size)
next(iter(data_iter)) # 用iter构造迭代器，next读取第一项

# 定义模型，使用nn库，nn是神经网络的缩写
net = nn.Sequential(nn.Linear(2, 1)) # 在PyTorch中，全连接层在Linear类中定义,第一个指 定输入特征形状，即2，第二个指定输出特征形状，输出特征形状为单个标量，因此为1

# 初始化模型参数，使用net[0].weight.data进行索引
net[0].weight.data.normal_(0, 0.01)
net[0].bias.data.fill_(0) # 过net[0]选择网络中的第一个图层，然后使用weight.data和bias.data方法访问参数。我们还可以使用替换方法normal_和fill_来重写参数值

# 定义损失函数均方损失
loss = nn.MSELoss()

# 定义优化算法，实例化一个SGD实例时，我们要指定优化的参数（可通过net.parameters()从我们的模型中获得） 以及优化算法所需的超参数字典。
trainer = torch.optim.SGD(net.parameters(), lr=0.03) # 小批量随机梯度下降只需要设置lr值，这里设置为0.03

# 训练模型，：在每个迭代周期里，我们将完整遍历一次数据集（train_data），不停地从中获取一个小批量的输 入和相应的标签。对于每一个小批量，我们会进行以下步骤:
# 1、通过调用net(X)生成预测并计算损失l（前向传播）。
# 2、通过进行反向传播来计算梯度。
# 3、通过调用优化器来更新模型参数。
num_epochs = 5
for epoch in range(num_epochs): # 注意整数循环得加上range()
    for X, y in data_iter:
        l = loss(net(X), y)
        trainer.zero_grad() # 优化器梯度提前清零
        l.backward()
        trainer.step() # 执行优化算法
    l = loss(net(features), labels) # 使用更新后的参数进行计算最终损失
    print(f'epoch:{epoch+1}  loss:{l:f}')

# 比较一下估计参数与实际的大小
w = net[0].weight.data
b = net[0].bias.data
print(f'w的误差：{true_w - w}')
print(f'b的误差：{true_b - b}')

'''
小结：
1、我们可以使用pytorch的高级API更简洁的实现模型
2、在pytorch中，data模块提供了数据处理工具，nn模块定义了大量的神经网络层和常见的损失函数
3、我们可以通过_结尾的方法快速直接将参数替换，从而初始化参数.
'''